import sys
import os
from pathlib import Path

from tests.base_test import base_test_case

logger = base_test_case.get_logger(__name__)
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir

try:
    import pypdfium2 as pdfium
    PYPDFIUM_AVAILABLE = True
except ImportError:
    PYPDFIUM_AVAILABLE = False
    logger.warning("pypdfium2 is not available. Please install it with: pip install pypdfium2")


def extract_text_with_coordinates(pdf_path: str, page_number: int = 0, scale: float = 1.0):
    """
    使用pypdfium2提取PDF页面的文本和坐标，支持缩放因子
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
        scale (float): 缩放因子
    
    Returns:
        list: 包含文本块和坐标的列表，每个元素为 (x0, y0, x1, y1, text)
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        # 获取指定页面
        page = doc[page_number]
        
        # 获取页面文本
        text_page = page.get_textpage()
        
        # 提取文本块和坐标
        text_blocks = []
        
        # 遍历文本块
        for i in range(text_page.count_rects()):
            rect = text_page.get_rect(i)
            x0, y0, x1, y1 = rect
            # 应用缩放因子到坐标
            x0, y0, x1, y1 = x0 * scale, y0 * scale, x1 * scale, y1 * scale
            text = text_page.get_text_bounded(*rect)
            if text.strip():
                text_blocks.append((x0, y0, x1, y1, text.strip()))
        
        return text_blocks
    finally:
        doc.close()


def extract_text_as_dict(pdf_path: str, page_number: int = 0, scale: float = 1.0):
    """
    使用pypdfium2提取PDF页面的详细文本信息，类似于PyMuPDF的page.get_text('dict')
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
        scale (float): 缩放因子
    
    Returns:
        dict: 包含详细文本信息的字典
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_number]
        text_page = page.get_textpage()
        
        # 构建类似于PyMuPDF的dict结构
        result = {
            "width": page.get_width() * scale,
            "height": page.get_height() * scale,
            "blocks": []
        }
        
        # 获取文本块信息
        for i in range(text_page.count_rects()):
            rect = text_page.get_rect(i)
            x0, y0, x1, y1 = rect
            # 应用缩放因子
            x0, y0, x1, y1 = x0 * scale, y0 * scale, x1 * scale, y1 * scale
            
            text = text_page.get_text_bounded(*rect)
            if text.strip():
                block = {
                    "bbox": [x0, y0, x1, y1],
                    "lines": [{
                        "bbox": [x0, y0, x1, y1],
                        "spans": [{
                            "bbox": [x0, y0, x1, y1],
                            "text": text.strip(),
                            "font": "unknown",  # pypdfium2不直接提供字体信息
                            "size": 0.0  # pypdfium2不直接提供字体大小
                        }]
                    }],
                    "type": 0  # 0表示文本块
                }
                result["blocks"].append(block)
        
        return result
    finally:
        doc.close()


def render_page_with_scale(pdf_path: str, page_number: int = 0, scale: float = 1.0):
    """
    使用指定缩放因子渲染PDF页面为图像
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
        scale (float): 缩放因子
    
    Returns:
        PIL.Image: 渲染的图像
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_number]
        # 使用指定缩放因子渲染页面
        image = page.render(scale=scale).to_pil()
        return image
    finally:
        doc.close()


def get_detailed_text_info(pdf_path: str, page_number: int = 0):
    """
    获取PDF页面的详细文本信息，包括字符级别的信息
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
    
    Returns:
        dict: 包含详细文本信息的字典
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_number]
        text_page = page.get_textpage()
        
        # 获取字符级别的信息
        char_count = text_page.count_chars()
        chars_info = []
        
        for i in range(char_count):
            # 获取字符边界框
            char_box = text_page.get_charbox(i)
            # 获取字符文本
            char_text = text_page.get_text_range(i, 1)
            
            chars_info.append({
                "text": char_text,
                "bbox": char_box
            })
        
        return {
            "char_count": char_count,
            "chars": chars_info
        }
    finally:
        doc.close()


def extract_images_info(pdf_path: str, page_number: int = 0):
    """
    提取PDF页面中的图像信息
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
    
    Returns:
        list: 包含图像信息的列表
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_number]
        
        # 获取页面中的图像对象
        image_objects = []
        for obj in page.get_objects():
            # 检查对象类型是否为图像
            if obj.type == pdfium.FPDF_PAGEOBJ_IMAGE:
                # 获取图像边界框
                bbox = obj.get_bounds()
                image_objects.append({
                    "bbox": bbox,
                    "type": "image"
                })
        
        return image_objects
    finally:
        doc.close()


def extract_tables_info(pdf_path: str, page_number: int = 0):
    """
    提取PDF页面中的表格信息（基于文本布局分析）
    
    Args:
        pdf_path (str): PDF文件路径
        page_number (int): 页码（从0开始）
    
    Returns:
        list: 包含表格信息的列表
    """
    if not PYPDFIUM_AVAILABLE:
        raise ImportError("pypdfium2 is not available. Please install it to use this feature.")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")
    
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_number]
        text_page = page.get_textpage()
        
        # 获取文本块并分析可能的表格结构
        text_blocks = []
        for i in range(text_page.count_rects()):
            rect = text_page.get_rect(i)
            text = text_page.get_text_bounded(*rect)
            if text.strip():
                text_blocks.append({
                    "bbox": rect,
                    "text": text.strip()
                })
        
        # 简单的表格检测逻辑（基于文本块的对齐）
        # 这是一个简化的实现，实际的表格检测需要更复杂的算法
        tables = []
        if len(text_blocks) > 4:  # 至少需要一些文本块才能形成表格
            # 按y坐标排序
            text_blocks.sort(key=lambda x: x["bbox"][1])
            
            # 简单的表格检测：查找具有相似y坐标且水平排列的文本块
            current_row = []
            last_y = None
            row_height_threshold = 5  # 行高阈值
            
            for block in text_blocks:
                x0, y0, x1, y1 = block["bbox"]
                if last_y is None or abs(y0 - last_y) < row_height_threshold:
                    current_row.append(block)
                else:
                    if len(current_row) >= 2:  # 至少两列才认为可能是表格
                        # 计算行的边界框
                        min_x = min(b["bbox"][0] for b in current_row)
                        max_x = max(b["bbox"][2] for b in current_row)
                        min_y = min(b["bbox"][1] for b in current_row)
                        max_y = max(b["bbox"][3] for b in current_row)
                        
                        tables.append({
                            "bbox": (min_x, min_y, max_x, max_y),
                            "rows": [current_row],
                            "type": "table"
                        })
                    current_row = [block]
                last_y = y0
            
            # 处理最后一行
            if len(current_row) >= 2:
                min_x = min(b["bbox"][0] for b in current_row)
                max_x = max(b["bbox"][2] for b in current_row)
                min_y = min(b["bbox"][1] for b in current_row)
                max_y = max(b["bbox"][3] for b in current_row)
                
                tables.append({
                    "bbox": (min_x, min_y, max_x, max_y),
                    "rows": [current_row],
                    "type": "table"
                })
        
        return tables
    finally:
        doc.close()


if __name__ == '__main__':
    """主测试函数"""
    if not PYPDFIUM_AVAILABLE:
        logger.error("Cannot run tests without pypdfium2. Please install it first.")
        sys.exit(1)
    
    # 查找测试PDF文件
    pdf_files = list(TEST_DATA_DIR.glob("*.pdf"))
    
    pdf_path = str(TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf")
    logger.info(f"Testing with PDF: {pdf_path}")
    
    # 测试不同缩放因子的文本提取
    scale_factors = [1.0, 2.0]
    
    for scale in scale_factors:
        logger.info(f"Testing text extraction with scale factor: {scale}")
        try:
            text_blocks = extract_text_with_coordinates(pdf_path, page_number=0, scale=scale)
            logger.info(f"Found {len(text_blocks)} text blocks")
            
            # 显示前几个文本块
            for i, (x0, y0, x1, y1, text) in enumerate(text_blocks[:3]):
                logger.info(f"  Block {i+1}: ({x0:.1f}, {y0:.1f}, {x1:.1f}, {y1:.1f}) - '{text[:50]}{'...' if len(text) > 50 else ''}")
        except Exception as e:
            logger.error(f"Error extracting text with scale {scale}: {e}")
    
    # 测试类似PyMuPDF的dict格式提取
    logger.info("Testing PyMuPDF-like dict extraction")
    try:
        dict_result = extract_text_as_dict(pdf_path, page_number=0, scale=1.0)
        logger.info(f"Dict result has {len(dict_result['blocks'])} blocks")
        
        # 显示前几个块的信息
        for i, block in enumerate(dict_result["blocks"][:2]):
            bbox = block["bbox"]
            text = block["lines"][0]["spans"][0]["text"]
            logger.info(f"  Block {i+1}: BBox=({bbox[0]:.1f}, {bbox[1]:.1f}, {bbox[2]:.1f}, {bbox[3]:.1f}), Text='{text[:50]}{'...' if len(text) > 50 else ''}'")
    except Exception as e:
        logger.error(f"Error in dict extraction: {e}")
    
    # 测试字符级别信息提取
    logger.info("Testing character-level information extraction")
    try:
        char_info = get_detailed_text_info(pdf_path, page_number=0)
        logger.info(f"Found {char_info['char_count']} characters")
        
        # 显示前几个字符的信息
        for i, char in enumerate(char_info["chars"][:5]):
            bbox = char["bbox"]
            text = char["text"]
            logger.info(f"  Char {i+1}: '{text}' at ({bbox[0]:.1f}, {bbox[1]:.1f}, {bbox[2]:.1f}, {bbox[3]:.1f})")
    except Exception as e:
        logger.error(f"Error in character-level extraction: {e}")
    
    # 测试图像提取
    logger.info("Testing image extraction")
    try:
        images = extract_images_info(pdf_path, page_number=0)
        logger.info(f"Found {len(images)} images")
        
        for i, image in enumerate(images[:3]):  # 显示前3个图像
            bbox = image["bbox"]
            logger.info(f"  Image {i+1}: BBox=({bbox[0]:.1f}, {bbox[1]:.1f}, {bbox[2]:.1f}, {bbox[3]:.1f})")
    except Exception as e:
        logger.error(f"Error in image extraction: {e}")
    
    # 测试表格提取
    logger.info("Testing table extraction")
    try:
        tables = extract_tables_info(pdf_path, page_number=0)
        logger.info(f"Found {len(tables)} potential tables")
        
        for i, table in enumerate(tables[:2]):  # 显示前2个表格
            bbox = table["bbox"]
            logger.info(f"  Table {i+1}: BBox=({bbox[0]:.1f}, {bbox[1]:.1f}, {bbox[2]:.1f}, {bbox[3]:.1f})")
    except Exception as e:
        logger.error(f"Error in table extraction: {e}")
    
    # 测试页面渲染
    logger.info("Testing page rendering")
    try:
        for scale in [1.0, 2.0]:
            image = render_page_with_scale(pdf_path, page_number=0, scale=scale)
            logger.info(f"Rendered page with scale {scale}: {image.width}x{image.height} pixels")
    except Exception as e:
        logger.error(f"Error rendering page: {e}")
